package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"sync"
	"time"

	pb "rpcperftest/perf"

	"google.golang.org/grpc"
)

var lock sync.Mutex

func testDownload(client pb.TestClient) {
	serverStream, err := client.DownloadRequest(context.Background(), &pb.DownloadOption{
		DurationTime: 120,
	})
	if err != nil {
		log.Fatalln("client.DownloadRequest err:", err)
	}
	var recvSize int64 = 0
	// report output
	go func() {
		interval := 2
		for {
			time.Sleep(time.Duration(interval) * time.Second)
			lock.Lock()
			fmt.Printf("[INFO] client recevie rate: %.2fMbps\n", float64(recvSize)/float64(interval)/1024/1024)
			recvSize = 0
			lock.Unlock()
		}
	}()
	// recv
	for {
		downloadData, err := serverStream.Recv()
		if err == io.EOF {
			recvSize += int64(downloadData.GetDataSize())
			break
		}
		if err != nil {
			log.Fatalln("serverStream.Recv err:", err)
		}
		lock.Lock()
		recvSize += int64(downloadData.GetDataSize())
		// fmt.Println("recv data size:", recvSize)
		lock.Unlock()
	}
}

func uploadRequest(client pb.TestClient) {
	clientStream, err := client.UploadRequest(context.Background())
	if err != nil {
		log.Fatalln("client.UploadRequest err:", err)
	}
	var testDataSize int32 = 4096
	uploadData := pb.UploadData{TestData: string(make([]byte, testDataSize)), DataSize: testDataSize}
	endTime := time.Now().Add(10 * time.Second) // 测试时长
	var sentSize int64 = 0
	// report output
	go func() {
		interval := 2
		for {
			time.Sleep(time.Duration(interval) * time.Second)
			lock.Lock()
			fmt.Printf("[INFO] client send rate: %.2fMbps\n", float64(sentSize)/float64(interval)/1024/1024)
			sentSize = 0
			lock.Unlock()
		}
	}()
	// send
	for {
		if err := clientStream.Send(&uploadData); err != nil {
			log.Fatalln("clientStream.Send err:", err)
		}
		lock.Lock()
		sentSize += int64(testDataSize)
		// fmt.Printf("%d byte send\n", sentSize)
		lock.Unlock()
		if time.Now().After(endTime) {
			break
		}
	}
	summary, err := clientStream.CloseAndRecv()
	if err != nil {
		log.Fatalln("clientStream.CloseAndRecv err:", err)
	}
	time.Sleep(time.Duration(2) * time.Second)
	fmt.Printf("[INFO] summary report: server received %d bytes, cost %d seconds, received rate: %.2fMbps\n", summary.BytesCount, summary.CostTime, summary.Rate)
}

func main() {
	conn, err := grpc.Dial("localhost:5000", grpc.WithInsecure(), grpc.WithBlock())
	if err != nil {
		log.Fatalln("grpc.Dial err:", err)
	}
	defer conn.Close()

	client := pb.NewTestClient(conn)
	testDownload(client)
	// uploadRequest(client)
}
